package com.fpf.handler;

import com.fpf.domain.GetAnswerDTO;
import com.fpf.service.IWindowTextService;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import java.io.IOException;
import java.util.function.Consumer;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;

public class ChatWebSocketHandler extends TextWebSocketHandler {
    private final IWindowTextService windowTextService;
    private static final ConcurrentHashMap<String, AtomicBoolean> stopFlags = new ConcurrentHashMap<>();

    public ChatWebSocketHandler(IWindowTextService windowTextService) {
        this.windowTextService = windowTextService;
    }


    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        System.out.println("客户端连接: " + session.getId());
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        // 处理停止请求
        if ("STOP".equals(payload)) {
            AtomicBoolean stopFlag = stopFlags.get(session.getId());
            if (stopFlag != null) {
                stopFlag.set(true);
            }
            return;
        }
        System.out.println("收到消息，开始流式处理: " + payload);
        // 创建一个用于接收分片的消费者
        GetAnswerDTO dto = new GetAnswerDTO();
        dto.setWindowContent(payload);

        AtomicBoolean stopFlag = new AtomicBoolean(false);//初始化一个停止标志，用于中断流式响应
        stopFlags.put(session.getId(), stopFlag);//将停止标志保存到map中

        Consumer<String> chunkConsumer = chunk -> {
            try {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage(chunk));
                }
            } catch (IOException e) {
                System.err.println("发送 WebSocket 消息失败: " + e.getMessage());
                throw new RuntimeException(e);
            }
        };
        windowTextService.getAnswerStream(dto, chunkConsumer, stopFlag);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, org.springframework.web.socket.CloseStatus status) {
        System.out.println("客户端断开: " + session.getId());
        stopFlags.remove(session.getId());
    }

}